/*
 *
 *  ------------------------------------------------------------------
 *  Copyright © 2017 Hangzhou DtDream Technology Co.,Lt d. All rights reserved.
 *  ------------------------------------------------------------------
 *        Product:  EMR
 *    Module Name: NEZHA
 *   Date Created: 17-10-23 下午4:22
 *    Description:
 *  ------------------------------------------------------------------
 *  Modification History
 *  DATE            Name           Description
 *  ------------------------------------------------------------------
 *  2017-10-23      NEZHA EMR
 *  ------------------------------------------------------------------
 * /
 */

package org.xukai.remoting.sdk.web.monitor;

import okhttp3.OkHttpClient;
import okhttp3.Request;
import okhttp3.Response;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.net.ssl.HostnameVerifier;
import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLSession;
import javax.net.ssl.SSLSocketFactory;
import javax.net.ssl.TrustManager;
import javax.net.ssl.X509TrustManager;
import java.io.IOException;
import java.io.Reader;
import java.security.KeyManagementException;
import java.security.NoSuchAlgorithmException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;
import org.xukai.remoting.sdk.web.AccessKey;
import org.xukai.remoting.sdk.web.api.RestAPIClient;
import org.xukai.remoting.sdk.web.exception.NeZhaClientException;

/**
 * Created by moriarty on 17-7-12.
 * <p>
 * Each Running SQL statement will be translated to several jobs by SparkThriftServer.
 * {@link SQLTracker} is a tool for tracking jobs of a SQL statement.
 * {@link SQLTracker} talks to Spark UI directly.
 */
public class SQLTracker {

    private static final Logger LOG = LoggerFactory.getLogger(SQLTracker.class);

    private static OkHttpClient httpClient;

    static {
        try {
            TrustManager[] trustAllCerts = new TrustManager[]{
                    new X509TrustManager() {
                        @Override
                        public void checkClientTrusted(java.security.cert.X509Certificate[] chain, String authType) {
                        }

                        @Override
                        public void checkServerTrusted(java.security.cert.X509Certificate[] chain, String authType) {
                        }

                        @Override
                        public java.security.cert.X509Certificate[] getAcceptedIssuers() {
                            return new java.security.cert.X509Certificate[]{};
                        }
                    }
            };
            SSLContext sslContext = SSLContext.getInstance("SSL");
            sslContext.init(null, trustAllCerts, new java.security.SecureRandom());
            SSLSocketFactory sslSocketFactory = sslContext.getSocketFactory();
            OkHttpClient.Builder builder = new OkHttpClient.Builder();
            httpClient = builder.sslSocketFactory(sslSocketFactory, (X509TrustManager) trustAllCerts[0])
                    .hostnameVerifier(new HostnameVerifier() {
                        @Override
                        public boolean verify(String hostname, SSLSession session) {
                            return true;
                        }
                    })
                    .readTimeout(1, TimeUnit.MINUTES)
                    .build();
        } catch (NoSuchAlgorithmException | KeyManagementException e) {
            throw new RuntimeException(e);
        }
    }

    enum RequestType {
        STATEMENT, JOB, STAGE
    }

    private AccessKey accessKey;
    private RestAPIClient restClient;
    private String database;
    private String jobId;
    private int stmtIndex;
    private boolean inited = false;
    private String currentUI;
    private String jobGroup;
    private Map<RequestType, String> requestUrlTemplate = new HashMap<>();

    public SQLTracker(AccessKey accessKey, RestAPIClient restClient, String database, String jobId, int stmtIndex) {
        this.accessKey = accessKey;
        this.restClient = restClient;
        this.database = database;
        this.jobId = jobId;
        this.stmtIndex = stmtIndex;
    }

    private boolean tryInit() throws IOException {
        if (inited) {
            return true;
        }

        try {
            SQLTrackerInfo trackerInfo = restClient.getJobStmtTrackerInfo(accessKey, database, jobId,
                    stmtIndex);
            this.currentUI = trackerInfo.getCurrentUI();
            String baseUrl = trackerInfo.getCurrentUI() + "/api/v1/applications/" + trackerInfo.getApplicationId();
            this.jobGroup = trackerInfo.getJobGroup();
            requestUrlTemplate.put(RequestType.STATEMENT, baseUrl + "/statements/%s");
            requestUrlTemplate.put(RequestType.JOB, baseUrl + "/jobs/%s");
            requestUrlTemplate.put(RequestType.STAGE, baseUrl + "/stages/%s");
            inited = true;
            return true;
        } catch (NeZhaClientException e) {
            if (e.getMessage().contains("Statement has not started yet.")) {
                LOG.info("Statement has not started yet.");
                return false;
            } else {
                throw new IOException(e.getMessage(), e);
            }
        }
    }

    static class JsonDeserializer<T> {

        private Class<T> clazz;

        public JsonDeserializer(Class<T> clazz) {
            this.clazz = clazz;
        }

        public T deserialize(Reader reader) throws IOException {
            return JacksonCodec.deserialize(reader, clazz);
        }
    }

    <T> T execute(RequestType type, JsonDeserializer<T> deserializer, Object... param) throws IOException {
        if (inited) {
            String url = String.format(requestUrlTemplate.get(type), param);
            try (Response response = httpClient.newCall(new Request.Builder().url(url).get().build()).execute()) {
                if (response.isSuccessful()) {
                    return deserializer.deserialize(response.body().charStream());
                } else {
                    throw new IOException(response.toString());
                }
            }
        } else {
            return null;
        }
    }

    /**
     * Gets Spark UI Address of this SQL.
     *
     * @return null if SQL has not started.
     * @throws IOException
     */
    public String getSparkUI() throws IOException {
        tryInit();
        return currentUI;
    }

    /**
     * Gets SQL statement state from Spark UI.
     *
     * @return
     * @throws IOException
     */
    public ExecutionState getExecutionState() throws IOException {
        StatementData statementData = getStatementData();
        if (statementData == null) {
            return ExecutionState.PENDDING;
        } else {
            return statementData.getState();
        }
    }

    /**
     * Gets latest data of all Jobs.
     *
     * @return null if SQL has not started.
     * @throws IOException
     */
    public List<JobData> getJobDatas() throws IOException {
        StatementData statementData = getStatementData();
        if (statementData == null) {
            return null;
        }

        List<JobData> jobDataList = new ArrayList<>(statementData.getJobId().length);
        for (String jobId : statementData.getJobId()) {
            jobDataList.add(getJobData(jobId));
        }

        return jobDataList;
    }

    /**
     * Gets SQL statement summary from Spark UI.
     *
     * @return null if SQL has not started.
     * @throws IOException
     */
    public StatementData getStatementData() throws IOException {
        tryInit();
        return execute(RequestType.STATEMENT, new JsonDeserializer<>(StatementData.class), jobGroup);
    }

    /**
     * Gets JobData by {@code jobId}.
     *
     * @param jobId
     * @return null if SQL has not started.
     * @throws IOException
     */

    public JobData getJobData(String jobId) throws IOException {
        tryInit();
        return execute(RequestType.JOB, new JsonDeserializer<>(JobData.class), jobId);
    }

    /**
     * Gets StageData List specified by {@code stageId}.
     * A stage may contains several attempts. Each StageData represents an attempt.
     *
     * @return
     */
    public List<StageData> getStageDatas(int stageId) throws IOException {
        tryInit();
        JsonDeserializer<List> deserializer = new JsonDeserializer<List>(List.class) {
            @Override
            public List<StageData> deserialize(Reader reader) throws IOException {
                return JacksonCodec.deserialize(reader, ArrayList.class, StageData.class);
            }
        };
        return execute(RequestType.STAGE, deserializer, stageId);
    }
}
